// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/benchmark.h>
#include "melon/init/Init.h"
#include <pollux/dwio/common/type_with_id.h>
#include <pollux/dwio/dwrf/writer/column_writer.h>
#include <pollux/dwio/dwrf/writer/writer.h>
#include <pollux/type/type.h>
#include <pollux/vector/flat_vector.h>

using namespace kumo::pollux::dwio::common;
using namespace kumo::pollux;
using namespace kumo::pollux::dwrf;

constexpr vector_size_t kVectorSize = 10000;
vector_size_t kNumIterations = 1000;

float genData(float pos) {
  return float(pos * static_cast<float>(3.14));
}

bool isNotNull(int32_t pos, int32_t null_every) {
  return pos == 0 || pos % null_every != 0;
}

void runBenchmark(int null_every) {
  melon::BenchmarkSuspender braces;

  auto type = CppToType<float>::create();
  auto typeWithId = TypeWithId::create(type, 1);
  auto pool = memory::memoryManager()->addLeafPool();
  VectorPtr vector;
  // Prepare input
  BufferPtr values = AlignedBuffer::allocate<float>(kVectorSize, pool.get());
  auto valuesPtr = values->asMutable<float>();

  BufferPtr nulls =
      AlignedBuffer::allocate<char>(bits::nbytes(kVectorSize), pool.get());
  auto* nullsPtr = nulls->asMutable<uint64_t>();

  uint32_t nullCount = 0;
  for (size_t i = 0; i < kVectorSize; ++i) {
    bool isPresent = isNotNull(i, null_every);
    bits::set_null(nullsPtr, i, !isPresent);
    if (isPresent) {
      valuesPtr[i] = genData(i);
    } else {
      ++nullCount;
    }
  }

  vector = std::make_shared<FlatVector<float>>(
      pool.get(),
      REAL(),
      nullCount == 0 ? nullptr : nulls,
      kVectorSize,
      values,
      std::vector<BufferPtr>{});

  // write
  braces.dismiss();

  for (auto i = 0; i < kNumIterations; i++) {
    auto config = std::make_shared<dwrf::Config>();
    WriterContext context{
        config,
        memory::memoryManager()->addRootPool("FloatColumnWriterBenchmark")};
    auto writer = BaseColumnWriter::create(context, *typeWithId, 0);
    writer->write(vector, common::Ranges::of(0, kVectorSize));
  }
}

BENCHMARK(FloatColumnWriterBenchmark2) {
  runBenchmark(2);
}

BENCHMARK(FloatColumnWriterBenchmark5) {
  runBenchmark(5);
}

BENCHMARK(FloatColumnWriterBenchmark10) {
  runBenchmark(10);
}

BENCHMARK(FloatColumnWriterBenchmark25) {
  runBenchmark(25);
}

BENCHMARK(FloatColumnWriterBenchmark50) {
  runBenchmark(50);
}

BENCHMARK(FloatColumnWriterBenchmark100) {
  runBenchmark(100);
}

BENCHMARK(FloatColumnWriterBenchmark500) {
  runBenchmark(500);
}

BENCHMARK(FloatColumnWriterBenchmark1000) {
  runBenchmark(1000);
}

BENCHMARK(FloatColumnWriterBenchmark2000) {
  runBenchmark(2000);
}

BENCHMARK(FloatColumnWriterBenchmark5000) {
  runBenchmark(5000);
}

BENCHMARK(FloatColumnWriterBenchmark10000) {
  runBenchmark(10000);
}

int32_t main(int32_t argc, char* argv[]) {
  melon::Init init{&argc, &argv};
  memory::MemoryManager::initialize({});
  melon::runBenchmarks();
  return 0;
}
